// readers/postgresql_reader.go
package readers

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"mig2kwdb/common"
	"reflect"
	"strings"

	_ "github.com/lib/pq"
)

type PostgreSQLReader struct {
	db          *sql.DB
	connInfo    common.DBConnectionInfo
	dataBases   []common.DatabaseConfig
	handler     common.ReplicateHandlerInterface
	replicators map[string]*PostgresqlReplicator
	ctx         context.Context
}

// NewPostgreSQLReader 创建新的 PostgreSQLReader 实例并初始化数据库连接
func NewPostgreSQLReader(connInfo common.DBConnectionInfo) (*PostgreSQLReader, error) {
	// 构建 PostgreSQL DSN
	// PostgreSQL 连接字符串格式: "host=localhost port=5432 user=username password=password dbname=mydb sslmode=disable"
	dsn := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
		connInfo.Ip, connInfo.Port, connInfo.User, connInfo.Password, connInfo.DatabaseName)
	//	var connectionString string = fmt.Sprintf("host=%s  port=%d user=%s password=%s dbname=%s sslmode=disable connect_timeout=3", HOST, PORT, USER, PASSWORD, DATABASE)

	if connInfo.SSL {
		dsn = fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=require",
			connInfo.Ip, connInfo.Port, connInfo.User, connInfo.Password, connInfo.DatabaseName)
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, err
	}

	return &PostgreSQLReader{
		db:       db,
		connInfo: connInfo,
	}, nil
}
func (r *PostgreSQLReader) GetObjectName(name string) string {
	escaped := strings.ReplaceAll(name, `"`, `""`)
	return `"` + escaped + `"`
}
func (r *PostgreSQLReader) IsNetworkError(err error) bool {
	etype := reflect.TypeOf(err)
	name := etype.String()
	if strings.EqualFold(name, "*escommon.InternelError") || strings.EqualFold(name, "pgx.PgError") || strings.EqualFold(name, "*pq.Error") {
		return false
	} else {
		return true
	}
}

func (r *PostgreSQLReader) GetStartPos() (common.ReplicateStartPos, error) {
	return common.ReplicateStartPos{}, nil
}
func (r *PostgreSQLReader) SetDataBases(dataBases []common.DatabaseConfig) {
	r.dataBases = dataBases
}
func (r *PostgreSQLReader) SetReplicateHandler(handler common.ReplicateHandlerInterface) {
	r.handler = handler
	if r.replicators == nil {
		r.replicators = make(map[string]*PostgresqlReplicator)
	}
	for _, database := range r.dataBases {
		replicator := r.replicators[database.Name]
		if replicator == nil {
			replicator = &PostgresqlReplicator{}
			r.replicators[database.Name] = replicator
		}
		tables := []string{}
		if database.Includes != nil {
			for _, table := range database.Includes {
				tables = append(tables, table.Name)
			}
		} else {
			tables, _ = r.GetDatabaseTables(database.Name, database.Schema)
		}
		replicator.Init(r.ctx, handler, database, r.connInfo, tables)
	}

}
func (r *PostgreSQLReader) StartReplicate(ctx context.Context, startPos common.ReplicateStartPos) error {
	for _, rep := range r.replicators {
		rep.Start(ctx)
	}
	return nil
}

// GetDB 返回 PostgreSQL 数据库连接
func (r *PostgreSQLReader) GetDB() *sql.DB {
	return r.db
}

// Close 关闭 PostgreSQL 数据库连接
func (r *PostgreSQLReader) Close() {
	if r.db != nil {
		r.db.Close()
	}
}
func (r *PostgreSQLReader) GetDataBaseType() string {
	return "postgresql"
}

// GetDataBases retrieves a list of all databases in the PostgreSQL instance
func (r *PostgreSQLReader) GetDataBases() ([]string, error) {
	query := `
		SELECT datname 
		FROM pg_database 
		WHERE datistemplate = false
		ORDER BY datname
	`

	rows, err := r.db.Query(query)
	if err != nil {
		log.Println(query, err)
		return nil, err
	}
	defer rows.Close()

	var databases []string
	for rows.Next() {
		var dbName string
		if err := rows.Scan(&dbName); err != nil {
			return nil, err
		}
		databases = append(databases, dbName)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return databases, nil
}

// GetDatabaseTables retrieves a list of all tables in the specified schema
func (r *PostgreSQLReader) GetDatabaseTables(databaseName, schemaName string) ([]string, error) {
	// 在 PostgreSQL 中，schema 更常用，如果未指定 schema，则使用 public
	schema := schemaName
	if schema == "" {
		schema = "public"
	}

	query := `
		SELECT tablename 
		FROM pg_tables 
		WHERE schemaname = $1
		ORDER BY tablename
	`

	rows, err := r.db.Query(query, schema)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tables []string
	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err != nil {
			return nil, err
		}
		tables = append(tables, tableName)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return tables, nil
}
func (r *PostgreSQLReader) MapDataTypes(column []common.KwdbColumn) []common.KwdbColumn {
	for i := range column {
		switch strings.ToUpper(column[i].Type) {
		case "BYTEA":
			column[i].Type = common.DataTypes.BLOB.Name
		case "REAL":
			column[i].Type = common.DataTypes.FLOAT8.Name
		case "BINARY":
			column[i].Type = common.DataTypes.CHAR.Name
		case "VARBINARY":
			column[i].Type = common.DataTypes.BLOB.Name
		case "YEAR":
			column[i].Type = common.DataTypes.INT2.Name
		default:
		}
	}
	return column

}
func (r *PostgreSQLReader) GetTableDefinition(databaseName, schemaName, tableName string) (common.TableDefinition, error) {
	var err error
	var def common.TableDefinition
	def.DatabaseName = databaseName
	def.SchemaName = schemaName
	def.TableName = tableName
	def.Columns, err = r.GetTableColumn(databaseName, schemaName, tableName)
	if err != nil {
		log.Println("GetTableColumn", err)
		return def, err
	}
	
	return def, nil
}

// GetTableIndexes retrieves index information for a specific table in a PostgreSQL database
func (r *PostgreSQLReader) GetTableIndexes(databaseName, schemaName, tableName string) ([]common.IndexInfo, error) {
	// 在 PostgreSQL 中，schema 更常用，如果未指定 schema，则使用 public
	schema := schemaName
	if schema == "" {
		schema = "public"
	}

	// 首先获取索引基本信息
	indexQuery := `
		SELECT 
			i.relname AS index_name,
			ix.indisunique AS is_unique,
			ix.indisprimary AS is_primary
		FROM pg_class t
		INNER JOIN pg_namespace ns ON t.relnamespace = ns.oid
		INNER JOIN pg_index ix ON t.oid = ix.indrelid
		INNER JOIN pg_class i ON ix.indexrelid = i.oid
		WHERE t.relname = $1
		  AND ns.nspname = $2
		  AND ix.indisvalid = true
		ORDER BY i.relname
	`

	indexRows, err := r.db.Query(indexQuery, tableName, schema)
	if err != nil {
		return nil, err
	}
	defer indexRows.Close()

	// 使用 map 来组织索引信息
	indexMap := make(map[string]*common.IndexInfo)

	// 先获取所有索引的基本信息
	for indexRows.Next() {
		var indexName string
		var isUnique bool
		var isPrimary bool

		err := indexRows.Scan(&indexName, &isUnique, &isPrimary)
		if err != nil {
			return nil, err
		}

		indexInfo := &common.IndexInfo{
			Name:      indexName,
			IsUnique:  isUnique,
			IsPrimary: isPrimary,
			Type:      "INDEX",
		}

		// 确定索引类型
		if isPrimary {
			indexInfo.Type = "PRIMARY"
		} else if isUnique {
			indexInfo.Type = "UNIQUE"
		}

		indexMap[indexName] = indexInfo
	}

	if err = indexRows.Err(); err != nil {
		return nil, err
	}

	// 然后获取索引的列信息
	columnQuery := `
		SELECT 
			i.relname AS index_name,
			a.attname AS column_name,
			ia.attnum AS column_position,
			(CASE WHEN ix.indoption[ia.attnum-1] & 1 = 1 THEN 'DESC' ELSE 'ASC' END) AS sort_order
		FROM pg_class t
		INNER JOIN pg_namespace ns ON t.relnamespace = ns.oid
		INNER JOIN pg_index ix ON t.oid = ix.indrelid
		INNER JOIN pg_class i ON ix.indexrelid = i.oid
		INNER JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
		INNER JOIN (SELECT generate_series(1, 32) AS attnum) ia ON ia.attnum <= array_length(ix.indkey, 1) AND a.attnum = ix.indkey[ia.attnum-1]
		WHERE t.relname = $1
		  AND ns.nspname = $2
		  AND ix.indisvalid = true
		ORDER BY i.relname, ia.attnum
	`

	columnRows, err := r.db.Query(columnQuery, tableName, schema)
	if err != nil {
		return nil, err
	}
	defer columnRows.Close()

	// 添加列信息到索引中
	for columnRows.Next() {
		var indexName string
		var columnName string
		var columnSeq int
		var sortOrder string

		err := columnRows.Scan(&indexName, &columnName, &columnSeq, &sortOrder)
		if err != nil {
			return nil, err
		}

		// 添加列信息到索引中
		indexColumn := common.IndexColumn{
			Name:      columnName,
			Position:  columnSeq,
			SortOrder: sortOrder,
		}

		if indexInfo, exists := indexMap[indexName]; exists {
			indexInfo.Columns = append(indexInfo.Columns, indexColumn)
		}
	}

	if err = columnRows.Err(); err != nil {
		return nil, err
	}

	// 将 map 转换为 slice
	var indexes []common.IndexInfo
	for _, indexInfo := range indexMap {
		indexes = append(indexes, *indexInfo)
	}

	return indexes, nil
}

// GetTableColumn retrieves column information for a specific table in a PostgreSQL database
func (r *PostgreSQLReader) GetTableColumn(databaseName, schemaName, tableName string) ([]common.KwdbColumn, error) {
	// 在 PostgreSQL 中，schema 更常用，如果未指定 schema，则使用 public
	schema := schemaName
	if schema == "" {
		schema = "public"
	}

	query := `
		SELECT 
			a.attname AS column_name,
			a.attnum AS ordinal_position,
			CASE 
				WHEN t.typcategory = 'A' THEN
					-- 处理数组类型
					replace(pg_catalog.format_type(a.atttypid, a.atttypmod), '[]', '') || '[]'
				ELSE
					-- 处理非数组类型
					pg_catalog.format_type(a.atttypid, a.atttypmod)
			END AS data_type,
			CASE 
				WHEN a.attlen > 0 THEN a.attlen
				WHEN a.atttypmod > 0 AND t.typcategory != 'A' THEN a.atttypmod - 4
				ELSE 0
			END AS character_maximum_length,
			COALESCE(
				CASE 
					WHEN t.typname = 'numeric' AND a.atttypmod > 0 THEN ((a.atttypmod - 4) >> 16) & 65535
					ELSE NULL
				END, 0
			) AS numeric_precision,
			COALESCE(
				CASE 
					WHEN t.typname = 'numeric' AND a.atttypmod > 0 THEN (a.atttypmod - 4) & 65535
					ELSE NULL
				END, 0
			) AS numeric_scale,
			CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
			COALESCE(pg_get_expr(d.adbin, d.adrelid), '') AS column_default,
			CASE WHEN EXISTS (
				SELECT 1 
				FROM pg_constraint c 
				WHERE c.conrelid = a.attrelid 
				  AND a.attnum = ANY(c.conkey) 
				  AND c.contype = 'p'
			) THEN 'YES' ELSE 'NO' END AS primary_key,
			CASE WHEN EXISTS (
				SELECT 1 
				FROM pg_constraint c 
				WHERE c.conrelid = a.attrelid 
				  AND a.attnum = ANY(c.conkey) 
				  AND c.contype = 'u'
			) THEN 'YES' ELSE 'NO' END AS unique_key,
			COALESCE(col_description(a.attrelid, a.attnum), '') AS column_comment,
			CASE 
				WHEN pg_get_expr(d.adbin, d.adrelid) LIKE 'nextval%' THEN 'YES' 
				ELSE 'NO' 
			END AS auto_increment,
			COALESCE(ic.character_set_name, '') AS character_set_name,
			COALESCE(cl.collname, '') AS collation_name
		FROM pg_attribute a
		JOIN pg_type t ON a.atttypid = t.oid
		JOIN pg_class c ON a.attrelid = c.oid
		JOIN pg_namespace n ON c.relnamespace = n.oid
		LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
		LEFT JOIN pg_collation cl ON a.attcollation = cl.oid
		LEFT JOIN information_schema.columns ic ON ic.table_schema = n.nspname 
			AND ic.table_name = c.relname 
			AND ic.column_name = a.attname
		WHERE a.attnum > 0 
		  AND NOT a.attisdropped		  
		  AND ic.table_catalog = $1
		  AND c.relname = $2
		  AND n.nspname = $3
		ORDER BY a.attnum
	`

	rows, err := r.db.Query(query, databaseName, tableName, schema)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var columns []common.KwdbColumn
	for rows.Next() {
		var col common.KwdbColumn
		err := rows.Scan(
			&col.Name,
			&col.Position,
			&col.Type,
			&col.Length,
			&col.Precision,
			&col.Scale,
			&col.Nullable,
			&col.DefaultValue,
			&col.PrimaryKey,
			&col.Unique,
			&col.Comment,
			&col.AutoIncrement,
			&col.Charset,
			&col.Collation,
		)
		if err != nil {
			return nil, err
		}
		col.DefaultValue.Valid = false
		col.DefaultValue.String = ""
		columns = append(columns, col)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return r.MapDataTypes(columns), nil
}
func (r *PostgreSQLReader) GetConditionRows(databaseName, schemaName, tableName string, columns []common.KwdbColumn, condition string) (common.RowsInterface, error) {
	orderby := ""
	keyNum := 0
	sql := "select "
	for idx, col := range columns {
		if idx > 0 {
			sql += ","
		}
		sql += " " + r.GetObjectName(col.Name)
		if col.PrimaryKey.String == "YES" {
			if keyNum > 0 {
				orderby += ","
			}
			orderby += " " + r.GetObjectName(col.Name)
			keyNum++
		}
	}
	//如果没有主键，则使用唯一索引
	if orderby == "" {
		keyNum = 0
		for _, col := range columns {
			if col.Unique.String == "YES" {
				if keyNum > 0 {
					orderby += ","
				}
				orderby += " " + r.GetObjectName(col.Name)
				keyNum++
			}
		}
	}
	sql += " from " + r.GetObjectName(schemaName) + "." + r.GetObjectName(tableName)
	if condition != "" {
		sql += " where " + condition
	}
	if orderby != "" {
		sql += " order by " + orderby
	}

	rows, err := r.db.Query(sql)
	return common.WrapSQLRows(rows), err
}
